import arrow
from app.db.models.users import User
from app.db.models.jobs import Job
from app.db.models.packages import Packages
from app.db.models.assets import Asset
from app.db.models.balances import Balances
from app.api.v2.jobs import asset_schedule, ping, http, content, security_event, ssl, vul
from flask import g


class OpenMonitorOrAlert:

    def get_user_info(self, username_list):
        user_ids = []
        package_ids = set()
        balance_ids = set()
        user_maps = dict()
        user_package_maps = dict()
        user_balance_maps = dict()
        users = User.objects.find({"username": {"$in": username_list}, "identified": True, "planId": {"$exists": True}})
        if users.count() == 0:
            return [], [], []

        for user in users:
            user_id = user.id
            package_id = user.planId
            balance_id = user.balanceId
            package_ids.add(package_id)
            balance_ids.add(balance_id)
            user_ids.append(user_id)
            user_maps[user_id] = user
            user_package_maps[user_id] = package_id
            user_balance_maps[user_id] = balance_id

        return list(package_ids), user_ids, user_maps, user_package_maps, list(balance_ids), user_balance_maps

    def get_job_info(self, user_ids, job_url_list, enable_alert=None, enable_monitor=None):
        query_dict = {"uid": {"$in": user_ids}, "verification.isVerified": True}

        if enable_alert is not None:
            query_dict['enableAlert'] = enable_alert
        if enable_monitor is not None:
            query_dict['enableMonitor'] = enable_monitor

        if job_url_list:
            query_dict.update({"targetUrl": {"$in": job_url_list}})
            jobs = Job.objects.find(query_dict)
        else:
            jobs = Job.objects.find(query_dict)

        if jobs.count() == 0:
            return []

        return jobs

    def get_package_info(self, package_ids):
        package_maps = dict()
        packages = Packages.objects.filter(pk__in=package_ids)
        for package in packages:
            package_maps[package.id] = package

        return package_maps

    def get_balance_info(self, balance_ids, switch):
        balance_maps = dict()
        arrow_now = arrow.utcnow()
        query_dict = {"_id": {"$in": balance_ids}}
        if switch:
            query_dict.update({'planExpireTime': {'$gt': arrow_now.datetime}})
        balances = Balances.objects.find(query_dict)
        for balance in balances:
            balance_maps[balance.id] = balance

        return balance_maps

    def get_asset_info(self, asset_ids):
        asset_maps = dict()
        assets = Asset.objects.filter(pk__in=asset_ids)
        for asset in assets:
            asset_maps[asset.id] = asset

        return asset_maps

    def distribute_vul(self, job_id, job, enable_monitor, custom_monitor_page, pre_login):
        # 下发Web漏洞监测 vul
        data = {
            'immediateExec': False,
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'vulSettings': job.vulSettings.to_dict(),
            'customMonitorPage': custom_monitor_page,
        }
        if pre_login:
            data['prelogin'] = pre_login
        vul(data=data)

    def distribute_ssl(self, job_id, job, enable_monitor, custom_monitor_page):
        # 下发SSL安全监测 ssl
        data = {
            'immediateExec': False,
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'sslSettings': job.sslSettings.to_dict(),
            'customMonitorPage': custom_monitor_page,
        }
        ssl(data=data)

    def distribute_security_event(self, job_id, job, enable_monitor, custom_monitor_page, pre_login):
        # 下发安全事件监测 security_event
        data = {
            'immediateExec': False,
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'securityEventSettings': job.securityEventSettings.to_dict(),
            'customMonitorPage': custom_monitor_page,
        }
        if pre_login:
            data['prelogin'] = pre_login
        security_event(data=data)

    def distribute_content(self, job_id, job, enable_monitor, custom_monitor_page, pre_login):
        # 下发内容合格监测 content
        data = {
            'immediateExec': False,
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'contentSettings': job.contentSettings.to_dict(),
            'customMonitorPage': custom_monitor_page,
        }
        if pre_login:
            data['prelogin'] = pre_login
        content(data)

    def distribute_http(self, job_id, job, enable_monitor, custom_monitor_page):
        # 下发HTTP监测 http
        http_settings = []
        for item in job.httpSettings:
            http_settings.append(item.to_dict())
        data = {
            'immediateExec': False,
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'httpSettings': http_settings,
            'customMonitorPage': custom_monitor_page,
        }
        http(data=data)

    def distribute_ping(self, job_id, job, enable_monitor, custom_monitor_page):
        # 下发Ping监测 ping
        ping_settings = []
        for item in job.pingSettings:
            ping_settings.append(item.to_dict())
        data = {
            'immediateExec': False,
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'pingSettings': ping_settings,
            'customMonitorPage': custom_monitor_page,
        }
        ping(data)

    def distribute_asset(self, job_id, job, enable_monitor, asset):
        # 下发资产变动监测 asset
        data = {
            'jobId': job_id,
            'sourceIp': job.sourceIp,
            'enableMonitor': enable_monitor,
            'taskSettings': asset.taskSettings.to_dict(),
        }
        asset_schedule(data=data)

    def distribute_monitor_config(self, asset, job, enable_monitor=True):
        # 下发配置任务到引擎
        job_id = str(job.id)

        custom_monitor_page = {'includeUrl': [], 'excludeUrl': []}
        if hasattr(job, 'customMonitorPage'):
            custom_monitor_page = job.customMonitorPage.to_dict()

        pre_login = dict()
        if hasattr(job, 'prelogin'):
            pre_login = job.prelogin.to_dict()

        # 下发Web漏洞监测 vul
        self.distribute_vul(job_id, job, enable_monitor, custom_monitor_page, pre_login)

        # 下发SSL安全监测 ssl
        self.distribute_ssl(job_id, job, enable_monitor, custom_monitor_page)

        # 下发安全事件监测 security_event
        self.distribute_security_event(job_id, job, enable_monitor, custom_monitor_page, pre_login)

        # 下发内容合格监测 content
        self.distribute_content(job_id, job, enable_monitor, custom_monitor_page, pre_login)

        # 下发HTTP监测 http
        self.distribute_http(job_id, job, enable_monitor, custom_monitor_page)

        # 下发Ping监测 ping
        self.distribute_ping(job_id, job, enable_monitor, custom_monitor_page)

        # 下发资产变动监测 asset
        self.distribute_asset(job_id, job, enable_monitor, asset)

    def check_has_number(self, task):
        if task.used >= task.total:
            return False
        return True

    def check_job_balance(self, balance):
        task_check_list = [balance.domainEnable, ]
        for task in task_check_list:
            if not self.check_has_number(task):
                return False
        return True

    def update_jobs(self, action, job_ids, enable_alert, enable_monitor):
        if action == 'monitor':
            Job.objects(pk__in=job_ids).update(enableMonitor=enable_monitor)
        else:
            Job.objects(pk__in=job_ids).update(enableAlert=enable_alert)

    def check_job_settings_interval(self, balance, job, asset):
        def _check(job_settings, balance_settings):
            if job_settings.enable and job_settings.interval < balance_settings.total:
                job_settings.interval = balance_settings.total

        _check(job.vulSettings, balance.vulMaxFreq)
        _check(job.securityEventSettings, balance.securityEventMaxFreq)
        _check(job.contentSettings, balance.contentMaxFreq)
        http_settings = job.httpSettings
        for item in http_settings:
            _check(item, balance.httpMaxFreq)

        ping_settings = job.pingSettings
        for item in ping_settings:
            _check(item, balance.pingMaxFreq)

        _check(job.sslSettings, balance.sslMaxFreq)
        _check(asset.taskSettings, balance.assetMaxFreq)

    def check_http_max_task_enable_used(self, http_settings, switch):
        used = 0
        for item in http_settings:
            if item.enable:
                if switch:
                    used += 1
                else:
                    used -= 1

        return used

    def check_ping_max_task_enable_used(self, ping_settings, switch):
        used = 0
        for item in ping_settings:
            if item.enable:
                if switch:
                    used += 1
                else:
                    used -= 1
        return used

    def check_site_port_trait_used(self, vul_settings, switch):
        used = 0
        if vul_settings.enableSitePortraitTriggerMonitor:
            if switch:
                used += 1
            else:
                used -= 1
        return used

    def check_site_domain_enable_used(self, switch):
        used = 0
        if switch:
            used += 1
        else:
            used -= 1
        return used

    def execute(self, action, usernames, job_urls, enable_alert: bool, enable_monitor: bool):
        if action not in ('monitor', 'alert'):
            return False, "操作类型不合法"

        if enable_monitor not in (True, False):
            return False, "enable_monitor参数错误: {}".format(enable_monitor)

        if enable_alert not in (True, False):
            return False, "enable_alert参数错误: {}".format(enable_alert)

        switch = False
        if (action == 'monitor' and enable_monitor) or (action == 'alert' and enable_alert):
            switch = True

        username_list = usernames.split(',')
        job_url_list = []
        success_job_list = []
        if job_urls is not None:
            job_url_list = job_urls.split(',')

        # 查询用户
        package_ids, user_ids, user_maps, user_package_maps, balance_ids, user_balance_maps = self.get_user_info(
            username_list)
        if not user_ids or not package_ids or not balance_ids:
            return False, "用户不存在或者没有开通套餐"

        # 根据用户筛选资产
        if action == 'monitor':
            if enable_monitor:
                _enable_monitor = False
            else:
                _enable_monitor = True
            jobs = self.get_job_info(user_ids, job_url_list, enable_monitor=_enable_monitor)
        else:
            if enable_alert:
                _enable_alert = False
            else:
                _enable_alert = True
            jobs = self.get_job_info(user_ids, job_url_list, enable_alert=_enable_alert)
        if not jobs:
            return False, "没有查询到符合要求的资产，有可能资产不存在、有可能资产已经在目标状态"

        # 查询用户套餐，过滤出有套餐的用户
        # package_maps = self.get_package_info(package_ids)
        # if not package_maps:
        #     return success_job_list

        # 查询用户套餐使用情况
        balance_maps = self.get_balance_info(balance_ids, switch)
        if not balance_maps:
            return False, "没有开通套餐或者套餐已过期"

        asset_ids = [job.assetId for job in jobs]
        asset_maps = self.get_asset_info(asset_ids)

        job_ids = []
        for job in jobs:
            uid = job.uid
            asset_id = job.assetId
            package_id = user_package_maps.get(uid, '')
            balance_id = user_balance_maps.get(uid, '')
            if not package_id or not balance_id:
                continue

            asset = asset_maps.get(asset_id)
            if asset is None:
                continue

            # 套餐和套餐使用情况
            # package = package_maps.get(package_id)
            balance = balance_maps.get(balance_id)
            if balance is None:
                continue
            if action == 'monitor' and enable_monitor:
                if not self.check_job_balance(balance.balance):
                    continue

            g.entity_type = 'user'
            g.user = user_maps[uid]
            g.role = user_maps[uid].roleIds[0]

            if action == 'alert':
                enable_monitor = job.enableMonitor

            self.check_job_settings_interval(balance.balance, job, asset)
            try:
                self.distribute_monitor_config(asset, job, enable_monitor)
            except Exception:
                continue

            if action == 'monitor':
                domain_enable_used_inc = self.check_site_domain_enable_used(switch)
                http_enable_used_inc = self.check_http_max_task_enable_used(job.httpSettings, switch)
                ping_enable_used_inc = self.check_ping_max_task_enable_used(job.pingSettings, switch)
                site_port_trait_used_inc = self.check_site_port_trait_used(job.vulSettings, switch)
                balance.balance.domainEnable.used += domain_enable_used_inc

                Balances.objects(pk=balance.id).update_one(inc__balance__domainEnable__used=domain_enable_used_inc,
                                                           inc__balance__httpMaxTaskEnable__used=http_enable_used_inc,
                                                           inc__balance__pingMaxTaskEnable__used=ping_enable_used_inc,
                                                           inc__balance__sitePortrait__used=site_port_trait_used_inc)

            success_job_list.append(job.targetUrl)
            job_ids.append(job.id)

        self.update_jobs(action, job_ids, enable_alert, enable_monitor)
        return True, success_job_list
